from typing import Tuple, Dict
import torch
from tqdm import tqdm

from dataset import Dataset
from config import device
import mytorch
from .log import Logger


class Runner(object):
    def __init__(
            self,
            model: mytorch.net.Module,
            loss: mytorch.loss.Loss,
            optimizer: mytorch.optimizer.Optimizer,
            scheduler: mytorch.scheduler.Scheduler,
            logger: Logger,
    ):
        self.model = model
        self.loss = loss
        self.optimizer = optimizer
        self.scheduler = scheduler

        self.model.to(device)
        self.loss.to(device)
        self.loss.eval()

        self.logger = logger

    def lr(self) -> Tuple[float]:
        return tuple(map(lambda p: float(p['lr']), self.optimizer.param_groups))

    def train_seg(self, dataset: Dataset, epoch: int) -> None:
        # 每代开始前调整学习率
        self.scheduler.step(epoch)
        # 先训练分割
        dataset.mode('train_seg')
        dataset.sample()
        self.model.mode('train_seg')
        # 记录 loss、tp、fn、fp、tn 信息
        calculate = {'count': 0, 'loss': [], 'loss_seg': [], 'loss_aux': [], 'tp': torch.zeros(4), 'pr': torch.zeros(4), 'gt': torch.zeros(4)}
        # 开始训练
        with tqdm(dataset) as iterable:
            for images, labels_seg in iterable:
                self.optimizer.zero_grad()
                seg, aux = self.model(images)
                # 损失函数
                loss_seg = self.loss(seg, labels_seg.to(device))
                loss_aux = self.loss(aux, labels_seg.to(device))
                loss = loss_seg + loss_aux
                loss.backward()
                self.optimizer.step()
                iterable.set_postfix(ordered_dict={'seg': loss_seg.item(), 'aux': loss_seg.item(), 'loss': loss_seg.item()})
                seg, aux, lab = seg.detach().cpu(), aux.detach().cpu(), labels_seg.detach().cpu()
                calculate['count'] += 1
                # loss 记录
                calculate['loss'].append(loss.item())
                calculate['loss_seg'].append(loss_seg.item())
                calculate['loss_aux'].append(loss_aux.item())
                # metric 记录
                seg = torch.eye(5)[seg.argmax(dim=1)].permute(0, 3, 1, 2)[:, 1:, :, :].type(torch.bool)
                lab = lab[:, 1:, :, :].type(torch.bool)
                calculate['tp'] += (seg & lab).sum(dim=(0, 2, 3)).detach()
                calculate['pr'] += seg.sum(dim=(0, 2, 3)).detach()
                calculate['gt'] += lab.sum(dim=(0, 2, 3)).detach()
                del images, labels_seg, seg, aux, loss_seg, loss_aux, loss
        # 日志计算
        losses = {
            'loss': sum(calculate['loss']) / (calculate['count'] + 1e-17),
            'loss_seg': sum(calculate['loss_seg']) / (len(calculate['loss_seg']) + 1e-17),
            'loss_aux': sum(calculate['loss_aux']) / (len(calculate['loss_aux']) + 1e-17),
        }
        dices: torch.Tensor = 2 * calculate['tp'] / (calculate['pr'] + calculate['gt'] + 1e-17)
        jaccard: torch.Tensor = calculate['tp'] / (calculate['pr'] + calculate['gt'] - calculate['tp'] + 1e-17)
        # liver, spleen, left_kidney, right_kidney
        metrics = {
            'dice': dices.sum().item() / (len(dices) + 1e-17),
            'jaccard': jaccard.sum().item() / (len(jaccard) + 1e-17),
        }
        metrics['score_seg'] = 0.5 * metrics['dice'] + 0.5 * metrics['jaccard']
        organs = {
            'liver': dices[0].item(),
            'spleen': dices[1].item(),
            'lkidney': dices[2].item(),
            'rkidney': dices[3].item(),
        }
        self.logger.log_epoch(prefix='seg', items={**losses, **metrics, **organs}, epoch=epoch)

    def train_cls(self, dataset: Dataset, epoch: int) -> None:
        # 每代开始前调整学习率
        self.scheduler.step(epoch)
        # 再训练分类
        dataset.mode('train_cls')
        dataset.sample()
        self.model.mode('train_cls')
        # 记录 loss、tp、fn、fp、tn 信息 -> A 类：患者水平统计，阳性只计算一次， B 类： 器官水平统计，阳性计算多次， C 类： 我自己统计每个器官的分类情况
        calculate = {
            'count': 0,
            'loss': [],
            'A': {'tp': 0, 'fn': 0, 'fp': 0, 'tn': 0},
            'B': {'tp': 0, 'fn': 0, 'fp': 0, 'tn': 0},
            'C': {'liver': 0, 'spleen': 0, 'lkidney': 0, 'rkidney': 0}
        }
        # 开始训练
        with tqdm(dataset) as iterable:
            for images, labels_cls in iterable:
                self.optimizer.zero_grad()
                cls = self.model(images)
                # 损失函数
                loss = self.loss(cls, labels_cls.to(device))
                loss.backward()
                self.optimizer.step()
                cls = cls.detach().cpu()
                iterable.set_postfix(ordered_dict={'loss': loss.item()})
                # loss 记录
                calculate['loss'].append(loss.item())
                # metric 记录
                cls = (cls > 0.5).type(torch.bool)
                lab = labels_cls.type(torch.bool)
                calculate['count'] += 1
                # A 类统计
                calculate['A']['tp'] += (cls.any() & lab.any()).item()
                calculate['A']['fn'] += (~cls.all() & lab.any()).item()
                calculate['A']['fp'] += (cls.any() & ~lab.all()).item()
                calculate['A']['tn'] += (~cls.all() & ~lab.all()).item()
                # B 类统计
                calculate['B']['tp'] += (cls & lab).sum().item()
                calculate['B']['fn'] += (~cls & lab).sum().item()
                calculate['B']['fp'] += (cls & ~lab).sum().item()
                calculate['B']['tn'] += (~cls & ~lab).sum().item()
                # C 类统计
                calculate['C']['liver'] += cls[0, 0] == lab[0, 0]
                calculate['C']['spleen'] += cls[0, 1] == lab[0, 1]
                calculate['C']['lkidney'] += cls[0, 2] == lab[0, 2]
                calculate['C']['rkidney'] += cls[0, 3] == lab[0, 3]
                del images, labels_cls, cls, loss
        # 日志计算
        losses = {
            'loss': sum(calculate['loss']) / (calculate['count'] + 1e-17),
        }
        metrics: Dict[str, float] = {
            'recall_A': calculate['A']['tp'] / (calculate['A']['tp'] + calculate['A']['fn'] + 1e-17),
            'f1_A': 2 * calculate['A']['tp'] / (2 * calculate['A']['tp'] + calculate['A']['fp'] + calculate['A']['fn'] + 1e-17),
            'acc_A': (calculate['A']['tp'] + calculate['A']['tn']) / (calculate['count'] + 1e-17),
            'acc_B': (calculate['B']['tp'] + calculate['B']['tn']) / (calculate['count'] * 4 + 1e-17),
            'f1_B': 2 * calculate['B']['tp'] / (2 * calculate['B']['tp'] + calculate['B']['fp'] + calculate['B']['fn'] + 1e-17),
        }
        metrics['score_cls'] = 0.3 * metrics['recall_A'] + 0.2 * metrics['f1_A'] + 0.2 * metrics['acc_B'] + 0.2 * metrics['f1_B'] + 0.1 * metrics['acc_A']
        # liver, spleen, left_kidney, right_kidney
        organs = {
            'liver': calculate['C']['liver'] / (calculate['count'] + 1e-17),
            'spleen': calculate['C']['spleen'] / (calculate['count'] + 1e-17),
            'lkidney': calculate['C']['lkidney'] / (calculate['count'] + 1e-17),
            'rkidney': calculate['C']['rkidney'] / (calculate['count'] + 1e-17),
        }
        self.logger.log_epoch(prefix='cls', items={**losses, **metrics, **organs}, epoch=epoch)

    def valid(self, dataset: Dataset, prefix: str, epoch: int) -> None:
        # 确保模型处在预测状态
        dataset.mode('evaluate')
        dataset.sample()
        self.model.mode('predict')
        with torch.no_grad():
            # 记录 loss、tp、fn、fp、tn 信息 -> A 类：患者水平统计，阳性只计算一次， B 类： 器官水平统计，阳性计算多次， C 类： 我自己统计每个器官的分类情况
            calculate = {
                'count': 0,
                'A': {'tp': 0, 'fn': 0, 'fp': 0, 'tn': 0},
                'B': {'tp': 0, 'fn': 0, 'fp': 0, 'tn': 0},
                'C': {'liver': 0, 'spleen': 0, 'lkidney': 0, 'rkidney': 0},
                'D': {'tp': torch.zeros(4), 'pr': torch.zeros(4), 'gt': torch.zeros(4)}
            }
            # 开始验证
            with tqdm(dataset) as iterable:
                for images, labels_seg, labels_cls in iterable:
                    seg, cls = self.model(images)
                    seg, cls = seg.detach(), cls.detach()
                    lab_seg, lab_cls = labels_seg, labels_cls
                    # metric 记录
                    # seg 记录
                    seg = torch.eye(5)[seg.argmax(dim=1)].permute(0, 3, 1, 2)[:, 1:, :, :].type(torch.bool)
                    lab = lab_seg[:, 1:, :, :].type(torch.bool)
                    calculate['count'] += 1
                    calculate['D']['tp'] += (seg & lab).sum(dim=(0, 2, 3)).detach()
                    calculate['D']['pr'] += seg.sum(dim=(0, 2, 3)).detach()
                    calculate['D']['gt'] += lab.sum(dim=(0, 2, 3)).detach()
                    # cls 记录
                    cls: torch.Tensor = (cls > 0.5).type(torch.bool)
                    lab: torch.Tensor = lab_cls.type(torch.bool)
                    # A 类统计
                    calculate['A']['tp'] += (cls.any() & lab.any()).item()
                    calculate['A']['fn'] += (~cls.all() & lab.any()).item()
                    calculate['A']['fp'] += (cls.any() & ~lab.all()).item()
                    calculate['A']['tn'] += (~cls.all() & ~lab.all()).item()
                    # B 类统计
                    calculate['B']['tp'] += (cls & lab).sum().item()
                    calculate['B']['fn'] += (~cls & lab).sum().item()
                    calculate['B']['fp'] += (cls & ~lab).sum().item()
                    calculate['B']['tn'] += (~cls & ~lab).sum().item()
                    # C 类统计
                    calculate['C']['liver'] += int(cls[0, 0] == lab[0, 0])
                    calculate['C']['spleen'] += int(cls[0, 1] == lab[0, 1])
                    calculate['C']['lkidney'] += int(cls[0, 2] == lab[0, 2])
                    calculate['C']['rkidney'] += int(cls[0, 3] == lab[0, 3])
                    del images, labels_seg, labels_cls, seg, cls, lab
            # 日志计算
            dices: torch.Tensor = 2 * calculate['D']['tp'] / (calculate['D']['pr'] + calculate['D']['gt'] + 1e-17)
            jaccard: torch.Tensor = calculate['D']['tp'] / (calculate['D']['pr'] + calculate['D']['gt'] - calculate['D']['tp'] + 1e-17)
            # liver, spleen, left_kidney, right_kidney
            metrics_seg = {
                'dice_D': dices.sum().item() / (len(dices) + 1e-17),
                'jaccard_D': jaccard.sum().item() / (len(jaccard) + 1e-17),
            }
            metrics_seg['score_seg'] = 0.5 * metrics_seg['dice_D'] + 0.5 * metrics_seg['jaccard_D']
            organs_seg = {
                'liver_D': dices[0].item(),
                'spleen_D': dices[1].item(),
                'lkidney_D': dices[2].item(),
                'rkidney_D': dices[3].item(),
            }
            metrics_cls: Dict[str, float] = {
                'recall_A': calculate['A']['tp'] / (calculate['A']['tp'] + calculate['A']['fn'] + 1e-17),
                'f1_A': 2 * calculate['A']['tp'] / (2 * calculate['A']['tp'] + calculate['A']['fp'] + calculate['A']['fn'] + 1e-17),
                'acc_A': (calculate['A']['tp'] + calculate['A']['tn']) / (calculate['count'] + 1e-17),
                'acc_B': (calculate['B']['tp'] + calculate['B']['tn']) / (calculate['count'] * 4 + 1e-17),
                'f1_B': 2 * calculate['B']['tp'] / (2 * calculate['B']['tp'] + calculate['B']['fp'] + calculate['B']['fn'] + 1e-17),
            }
            metrics_cls['score_cls'] = 0.3 * metrics_cls['recall_A'] + 0.2 * metrics_cls['f1_A'] + \
                                       0.2 * metrics_cls['acc_B'] + 0.2 * metrics_cls['f1_B'] + 0.1 * metrics_cls['acc_A']
            # liver, spleen, left_kidney, right_kidney
            organs_cls = {
                'liver_C': calculate['C']['liver'] / (calculate['count'] + 1e-17),
                'spleen_C': calculate['C']['spleen'] / (calculate['count'] + 1e-17),
                'lkidney_C': calculate['C']['lkidney'] / (calculate['count'] + 1e-17),
                'rkidney_C': calculate['C']['rkidney'] / (calculate['count'] + 1e-17),
            }
            score = 0.3 * metrics_seg['score_seg'] + 0.7 * metrics_cls['score_cls']
            self.logger.log_epoch(prefix=prefix, items={'score': score, **metrics_seg, **organs_seg, **metrics_cls, **organs_cls}, epoch=epoch)
